rm(list=ls())


# adjust path below
###############################

load("C:\\DeuchertHuberSchelker.RData")


# functions
#############################

effects<-function(y,z,d,t){
  p10=mean(d[z==0]); p11=mean(d[z==1]); p01=1-p11; p00=1-p10; comp=p11-p10 
  LATE=(mean(y[t==1 & z==1])-mean(y[t==1 & z==0]))/comp
  comp0=mean(y[t==1 & z==0 & d==0])-p01/comp*(mean(y[t==0 & z==1 & d==0])-mean(y[t==0 & z==0 & d==0]))
  comp1=(mean(y[t==1 & z==1]*d[t==1 & z==1])-mean(y[t==1 & z==0]*d[t==1 & z==0]))/comp
  totalcomp= comp1-comp0
  ate=mean(y[t==1 & z==1])-mean(y[t==1 & z==0])
  dir0=mean(y[t==1 & z==1 & d==0])-mean(y[t==0 & z==1 & d==0])-(mean(y[t==1 & z==0 & d==0])-mean(y[t==0 & z==0 & d==0]))
  indir1=totalcomp-dir0
  results=c(ate, dir0, totalcomp, dir0, indir1, LATE, indir1-LATE )
  results
}

bootstrap<-function(y,z,d,t,boot=1999, clusters=NULL, key=NULL){
  if (1*(is.null(clusters))==0) obs<-length(clusters)
  if (1*(is.null(clusters))==1) obs<-length(y)
  mc=c()
  i=1
  while(i<=boot){
    db<-c(); yb<-c(); zb<-c(); tb<-c() 
    if (1*(is.null(clusters))==0){
      sboot<-sample(clusters,obs,TRUE)
      for (k in 1:length(sboot)) {
        db<-c(db,d[key==sboot[k]]); zb<-c(zb,z[key==sboot[k]]); yb<-c(yb,y[key==sboot[k]]); tb<-c(tb,t[key==sboot[k]])
      }
    }
    if (1*(is.null(clusters))==1){
      sboot<-sample(1:obs,obs,TRUE)
      yb=y[sboot]; db<-d[sboot]; zb=z[sboot]; tb<-t[sboot];
    }
    ef=effects(y=yb,z=zb,d=db,t=tb)  
    mc<-rbind(mc, c(ef))
    i=i+1
  }
  mc
}

strata.mediation<-function(y,z,d,t,boot=1999,clusters=NULL, key=NULL){
  eff=effects(y,z,d,t)  
  mc=bootstrap(y,z,d,t, boot=boot, clusters=clusters, key=key)
  sd=c(sd(as.numeric(mc[,1])),sd(as.numeric(mc[,2])),sd(as.numeric(mc[,3])),sd(as.numeric(mc[,4])),sd(as.numeric(mc[,5])),sd(as.numeric(mc[,6])),sd(as.numeric(mc[,7])))
  results=rbind(eff, sd, 2*pnorm(-abs(eff/sd)))
  results
}


# Application
#############################################

# run estimations
attach(data)
after=wave>4
# remark: for placebo tests: after=wave>3
application<-function(yy){
  dd=treatment1
  zz=ceiling
  tt=after
  ind=(is.na(yy)==0 & is.na(dd)==0 & is.na(zz)==0 & wave>=4)
  # remark: for placebo tests: ind=(is.na(yy)==0 & is.na(dd)==0 & is.na(zz)==0 & wave>=3 & wave<=4)
  y=yy[ind==1]; z=zz[ind==1]; d=dd[ind==1]; t=tt[ind==1]
  
  key=id[ind==1]
  keytemp<-key; key<-sort(key)
  clusters<--9
  for (i in 1:length(key)){
    if (key[i]>max(clusters)){
      clusters=c(clusters,key[i])
    }
  }
  cluster<-clusters[2:length(clusters)]
  key<-keytemp 
  
  results=strata.mediation(y,z,d,t,boot=1999, clusters=cluster, key=key)
  results
}

outcomes=c("repu",  "srepu", "demo", "sdemo",   "polalient", "Vietnam", "ERindex")
se=c("se.repu", "se.srepu", "se.demo", "se.sdemo",   "se.polalient", "se.Vietnam", "se.ERindex")
pval=c("p.repu", "p.srepu", "p.demo", "p.sdemo",   "p.polalient", "p.Vietnam", "p.ERindex")
yy=cbind(repu, srepu,  demo, sdemo,  polalient, Vietnam, ERindex)

results=c()
for(j in (1:ncol(yy))){ 
  results1=application(yy[,j])
  rownames(results1)<-c(outcomes[j], se[j], pval[j])
  results=rbind(results,results1)
}
colnames(results)<-c("ate",  "d nt", "t co (A6)", "d co 0", "in co 1 (A6)", "LATE", "in co 1 (A6)-LATE"  )


# display results of Tables 2, 4, and A2
#########################

round(results, digits=3)
# remarks: the first column corresponds to Table 2 in the main text and the first column of Table A2 in the appendix
# columns 2-7 correspond to Table 4 in the main text and columns 2-7 of Table A2 in the appendix 



